import netron
import torch
import hiddenlayer as h


def visual_net(net, input=torch.ones([2, 1, 40, 40])):
    o = net(input)
    torch.onnx.export(net, input, "test.onnx", opset_version=11)
    netron.start("test.onnx")

def visual_netbyhiddenlayer(net,input=torch.ones([2, 1, 40, 40])):

    myNetGraph = h.build_graph(net, input)  # 建立网络模型图
    # myNetGraph.theme = h.graph.THEMES['blue']  # blue 和 basic 两种颜色，可以不要
    myNetGraph.save(path='./demoModel.png', format='png')  # 保存网络模型图，可以设置 png 和 pdf 等

# from .xrsy.basic import UnResNet18
# from .xrsy.basic_Att import UnResNet18_Att
# from .xrsy.basic_Att_MDS import BASIC_ATT_MDS
# from .xrsy.basic_Att_MDS_MFE import BASIC_ATT_MDS_MFE
# from .xrsy.basic_Att_MDS_MFE_IR import BASIC_ATT_MDS_MFE_IR
from model.models import UnResNet18_Att as Model
net=Model()
visual_net(net)
